package com.huan.es8.aggregations.bucket;

import co.elastic.clients.elasticsearch._types.InlineScript;
import co.elastic.clients.elasticsearch._types.ScriptLanguage;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.mapping.RuntimeFieldType;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.util.NamedValue;
import com.huan.es8.AbstractEs8Api;
import org.junit.jupiter.api.*;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;

/**
 * Terms 聚合
 * @see  <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations-bucket-terms-aggregation.html">官方文档</a>
 * @see  <a href="https://blog.csdn.net/fu_huo_1993/article/details/127822355">博客地址</a>
 * @author huan.fu
 * @date 2022/11/13 - 16:04
 */
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
public class TermsAggs extends AbstractEs8Api {

    @BeforeAll
    public void createIndex() throws IOException {
        createIndex("index_person",
                "{\n" +
                        "  \"settings\": {\n" +
                        "    \"number_of_shards\": 1\n" +
                        "  },\n" +
                        "  \"mappings\": {\n" +
                        "    \"properties\": {\n" +
                        "      \"id\": {\n" +
                        "        \"type\": \"long\"\n" +
                        "      },\n" +
                        "      \"name\": {\n" +
                        "        \"type\": \"keyword\"\n" +
                        "      },\n" +
                        "      \"sex\": {\n" +
                        "        \"type\": \"keyword\"\n" +
                        "      },\n" +
                        "      \"age\": {\n" +
                        "        \"type\": \"integer\"\n" +
                        "      },\n" +
                        "      \"province\": {\n" +
                        "        \"type\": \"keyword\"\n" +
                        "      },\n" +
                        "      \"address\": {\n" +
                        "        \"type\": \"text\",\n" +
                        "        \"analyzer\": \"ik_max_word\",\n" +
                        "        \"fields\": {\n" +
                        "          \"keyword\": {\n" +
                        "            \"type\": \"keyword\",\n" +
                        "            \"ignore_above\": 256\n" +
                        "          }\n" +
                        "        }\n" +
                        "      }\n" +
                        "    }\n" +
                        "  }\n" +
                        "}");

        bulk("index_person", Arrays.asList(
                "{\"id\":1,\"name\":\"张三\",\"sex\":\"男\",\"age\":20,\"province\":\"湖北\",\"address\":\"湖北省黄冈市罗田县匡河镇\"}",
                "{\"id\":2,\"name\":\"李四\",\"sex\":\"男\",\"age\":19,\"province\":\"江苏\",\"address\":\"江苏省南京市\"}",
                "{\"id\":3,\"name\":\"王武\",\"sex\":\"女\",\"age\":25,\"province\":\"湖北\",\"address\":\"湖北省武汉市江汉区\"}",
                "{\"id\":4,\"name\":\"赵六\",\"sex\":\"女\",\"age\":30,\"province\":\"北京\",\"address\":\"北京市东城区\"}",
                "{\"id\":5,\"name\":\"钱七\",\"sex\":\"女\",\"age\":16,\"province\":\"北京\",\"address\":\"北京市西城区\"}",
                "{\"id\":6,\"name\":\"王八\",\"sex\":\"女\",\"age\":45,\"province\":\"北京\",\"address\":\"北京市朝阳区\"}"
        ));
    }

    @Test
    @DisplayName("统计人数最多的2个省")
    public void agg01() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg -> agg.terms(terms -> terms.field("province").size(2)))
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("统计人数最少的2个省-不推荐此写法，推荐使用 rare_terms")
    public void agg02() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms ->
                                terms.field("province")
                                        .size(2)
                                        .order(new NamedValue<>("_count", SortOrder.Asc))
                        )
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("根据年龄聚合，返回年龄最小的2个聚合")
    public void agg03() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms ->
                                terms.field("age")
                                        .size(2)
                                        .order(new NamedValue<>("_key", SortOrder.Asc))
                        )
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("子聚合排序-先根据省聚合，然后根据每个聚合后的最小年龄排序")
    public void agg04() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms ->
                                        terms.field("province")
                                                .size(10)
                                                .order(new NamedValue<>("min_age", SortOrder.Asc))
                                )
                                .aggregations("min_age", minAgg ->
                                        minAgg.min(min -> min.field("age")))
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("脚本聚合-根据省聚合，如果地址中有黄冈市则需要出现黄冈市")
    public void agg05() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .runtimeMappings("province_dynamic", field -> {
                    field.type(RuntimeFieldType.Keyword);
                    field.script(script -> script.inline(new InlineScript.Builder()
                            .lang(ScriptLanguage.Painless)
                            .source("String province = doc['province'].value;\n" +
                                    "          String address = doc['address.keyword'].value;\n" +
                                    "          if(address.contains('黄冈市')){\n" +
                                    "            emit('黄冈市');\n" +
                                    "          }else{\n" +
                                    "            emit(province);\n" +
                                    "          }")
                            .build()));
                    return field;
                })
                .aggregations("agg_province", agg ->
                        agg.terms(terms ->
                                terms.field("province_dynamic")
                                        .size(10)
                        )
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("filter-以省分组，并且只包含北的省，但是需要排除湖北省")
    public void agg06() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms ->
                                terms.field("province")
                                        // 分桶的数据中 需要 存在 北 这个字的
                                        .include(include -> include.regexp(".*北.*"))
                                        // 分桶的数据中 不可出现 湖北。
                                        .exclude(exclude -> exclude.terms(Collections.singletonList("湖北")))
                                        .size(10)
                        )
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("多term聚合-根据省和性别聚合，然后根据最大年龄倒序")
    public void agg07() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("province_and_sex", agg ->
                        agg.multiTerms(multiTerms ->
                                        multiTerms.terms(term -> term.field("province"))
                                                .terms(term -> term.field("sex"))
                                                .order(new NamedValue<>("max_age", SortOrder.Desc))
                                )
                                .aggregations("max_age", ageAgg ->
                                        ageAgg.max(max -> max.field("age")))

                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @Test
    @DisplayName("missing value")
    public void agg08() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms -> terms.field("province")
                                // 如果数据中缺少province字段，则给默认值defaultProvince
                                .missing("defaultProvince"))
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }
    @Test
    @DisplayName("多个聚合-同时返回根据省聚合和根据性别聚合")
    public void agg09() throws IOException {

        SearchRequest searchRequest = new SearchRequest.Builder()
                .size(0)
                .index("index_person")
                .aggregations("agg_province", agg ->
                        agg.terms(terms -> terms.field("province"))
                )
                .aggregations("agg_sex", agg ->
                        agg.terms(terms -> terms.field("sex"))
                )
                .build();
        System.out.println(searchRequest);
        SearchResponse<Object> response = client.search(searchRequest, Object.class);
        System.out.println(response);
    }

    @AfterAll
    public void deleteIndex() throws IOException {
        deleteIndex("index_person");
    }
}
